package edu.arizona.sista.learn.activity.fluents.paths;

import java.util.Comparator;
import java.util.TreeMap;
import java.util.Iterator;
import java.util.Vector;

import edu.arizona.sista.learn.activity.model.ActivityInstance;
import edu.arizona.sista.learn.activity.model.Predicate;
import edu.arizona.sista.learn.activity.model.PredicateInterval;
import edu.arizona.sista.learn.activity.util.PredicateNames;

public class PredicateExtractor {
	
	public static enum DataType {
		Points,
		Heatmaps
	}
	
	protected Scene scene;
	protected DataType dataType;
	
	public PredicateExtractor(Scene scene, DataType type) {
		this.scene = scene;
		this.dataType = type;
	}
	
	public ActivityInstance getPredicateIntervals(Vector<Vector<Predicate>> predicates) {
		return getPredicateIntervals(predicates, true);
	}
	
	public ActivityInstance getPredicateIntervals(Vector<Vector<Predicate>> predicates, boolean smooth) {		
		TreeMap<Predicate, Vector<Integer>> framesPerPredicate = new TreeMap<Predicate, Vector<Integer>>(new Comparator<Predicate>() {
			public int compare(Predicate p1, Predicate p2) {
				return p1.toString().compareTo(p2.toString());
			}
		});
		
		for (int i = 0; i < predicates.size(); i++) {
			for (Predicate p : predicates.get(i)) {
				if (framesPerPredicate.containsKey(p)) {
					framesPerPredicate.get(p).add(i);
				}
				else {
					Vector<Integer> value = new Vector<Integer>();
					value.add(i);
					framesPerPredicate.put(p, value);
				}
			}
		}
		
		ActivityInstance fluents = new ActivityInstance(1 /* give it some id */ );
		
				
		for (Predicate p : framesPerPredicate.keySet()) {
			Vector<Integer> frames = framesPerPredicate.get(p);
			int start = -1;
			int end = -1;
			for (int i = 0; i < frames.size(); i++) {
				if (end == -1) {
					start = end = frames.get(i);
				}
				else if (frames.get(i) == end + 1) {
					end = frames.get(i);
				}
				else {
					fluents.addInterval(new PredicateInterval(p, start, end));
					start = end = -1;
					i--;
					continue;
				}
			}
			if (end != -1) {
				fluents.addInterval(new PredicateInterval(p, start, end));
			}
		}
		
		if (smooth)
			fluents.smoothFragmentedFluents();
		
		return fluents;
	}
	
	public Vector<Vector<Predicate>> extractPredicates() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>(scene.getNumFrames());
		for (int i = 0; i < scene.getNumFrames(); i++) {
			retval.add(new Vector<Predicate>());
		}
		
		Vector<Vector<Predicate>> attached = extractAttached();
		Vector<Vector<Predicate>> moving = extractMoving();
		Vector<Vector<Predicate>> movingFastHuman = extractMovingFastHuman();
		Vector<Vector<Predicate>> distanceDecreasing = extractDistanceDecreasing();
		Vector<Vector<Predicate>> distanceIncreasing = extractDistanceIncreasing();
		Vector<Vector<Predicate>> distanceDecreasingToStationary = extractDistanceDecreasingToStationary();
		Vector<Vector<Predicate>> distanceIncreasingToStationary = extractDistanceIncreasingToStationary();
		Vector<Vector<Predicate>> distanceStationary = extractDistanceStationary();
		Vector<Vector<Predicate>> trajectoryStable = extractTrajectoryStable();
		Vector<Vector<Predicate>> speedIncreasing = extractSpeedIncreasing();
		Vector<Vector<Predicate>> speedDecreasing = extractSpeedDecreasing();
		
		for (int i = 0; i < scene.getNumFrames(); i++) {
			retval.get(i).addAll(attached.get(i));
			retval.get(i).addAll(moving.get(i));
			retval.get(i).addAll(movingFastHuman.get(i));
			retval.get(i).addAll(distanceDecreasing.get(i));
			retval.get(i).addAll(distanceIncreasing.get(i));
			retval.get(i).addAll(distanceDecreasingToStationary.get(i));
			retval.get(i).addAll(distanceIncreasingToStationary.get(i));
			retval.get(i).addAll(distanceStationary.get(i));
			retval.get(i).addAll(trajectoryStable.get(i));
			retval.get(i).addAll(speedIncreasing.get(i));
			retval.get(i).addAll(speedDecreasing.get(i));
		}
		
		return retval;
	}
	
	public double getThreshold(String predicateName, SceneEntity entity) {
		double moving = 0.10;
		if (PredicateNames.DISTANCE_DECREASING.equalsIgnoreCase(predicateName)) {
			return moving;
		}
		else if (PredicateNames.DISTANCE_DECREASING_TO_STATIONARY.equalsIgnoreCase(predicateName)) {
			return moving;
		}
		else if (PredicateNames.DISTANCE_INCREASING.equalsIgnoreCase(predicateName)) {
			return moving;
		}
		else if (PredicateNames.DISTANCE_INCREASING_TO_STATIONARY.equalsIgnoreCase(predicateName)) {
			return moving;
		}
		else if (PredicateNames.DISTANCE_STABLE.equalsIgnoreCase(predicateName)) {
			return moving;
		}
		else if (PredicateNames.MOVING.equalsIgnoreCase(predicateName)) {
			return moving;
		}
		else if (PredicateNames.MOVING_FAST_HUMAN.equalsIgnoreCase(predicateName)) {
			return 0.5;		// best within [.3-.4] (we get 'chase' at .2-.3)
		}
		else if (PredicateNames.SPEED_INCREASING.equalsIgnoreCase(predicateName)) {
			return 0.0016;
		}
		else if (PredicateNames.SPEED_DECREASING.equalsIgnoreCase(predicateName)) {
			return 0.0016;
		}
		else if (PredicateNames.TRAJECTORY_STABLE.equalsIgnoreCase(predicateName)) {
			return 0.075;	// best within [.035-.05]
		}
		else {
			// why are we here?
			return 0;
		}
	}
	
	/*
	 * Extract all the ATTACHED predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractAttached() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames(); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e1 : scene.getEntities()) {
				SceneEntity e2 = e1.getAttachedEntity(i);
				if (e2 != null) {
					frame.add(new Predicate(PredicateNames.ATTACHED, e1.getPredArg(), e2.getPredArg()));
					frame.add(new Predicate(PredicateNames.ATTACHED, e2.getPredArg(), e1.getPredArg()));
				}
			}
			retval.add(frame);
		}
		
		return retval;
	}
	
	/*
	 * Extract all the MOVING predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractMoving() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e : scene.getEntities()) {
				Point3D p1 = e.getPoint(i);
				Point3D p2 = e.getPoint(i+1);
				
				if (p1 != null) {
					if (p2 == null) {
						// TODO maybe we shouldn't include this predicate?
						// the entity is moving if it exited the scene
						frame.add(new Predicate(PredicateNames.MOVING, e.getPredArg()));
					}
					else {
						if (p1.distanceTo(p2) >= getThreshold(PredicateNames.MOVING, e)) {
							// the entity is moving if the movement is greater than minDistance
							frame.add(new Predicate(PredicateNames.MOVING, e.getPredArg()));
						}
						// Need to be not-moving if in frame but not moving 8/16 JBW
						else {
							frame.add(new Predicate(PredicateNames.NOT_MOVING, e.getPredArg()));
						}
					}
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	/*
	 * Extract all the MOVING_FAST_HUMAN predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractMovingFastHuman() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e : scene.getEntities()) {
				Point3D p1 = e.getPoint(i);
				Point3D p2 = e.getPoint(i+1);
				
				if (p1 != null && p2 != null) {
					if (p1.distanceTo(p2) >= getThreshold(PredicateNames.MOVING_FAST_HUMAN, e)) {
						// the entity is moving if the movement is greater than minDistance
						frame.add(new Predicate(PredicateNames.MOVING_FAST_HUMAN, e.getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	/*
	 * Extract all the DISTANCE_DECREASING predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistanceDecreasing() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				Point3D p1e1 = pair[0].getPoint(i);
				Point3D p1e2 = pair[1].getPoint(i);
				Point3D p2e1 = pair[0].getPoint(i+1);
				Point3D p2e2 = pair[1].getPoint(i+1);
				
				if (p1e1 == null || p1e2 == null || p2e1 == null || p2e2 == null) {
					continue;
				}
				
				double d1 = p1e1.distanceTo(p1e2);
				double d2 = p2e1.distanceTo(p2e2);
				
				if (d1 - d2 >= getThreshold(PredicateNames.DISTANCE_DECREASING, pair[0])) {
					frame.add(new Predicate(PredicateNames.DISTANCE_DECREASING, pair[0].getPredArg(), pair[1].getPredArg()));
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	/*
	 * Extract all the DISTANCE_INCREASING predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistanceIncreasing() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				Point3D p1e1 = pair[0].getPoint(i);
				Point3D p1e2 = pair[1].getPoint(i);
				Point3D p2e1 = pair[0].getPoint(i+1);
				Point3D p2e2 = pair[1].getPoint(i+1);
				
				if (p1e1 == null || p1e2 == null || p2e1 == null || p2e2 == null) {
					continue;
				}
				
				double d1 = p1e1.distanceTo(p1e2);
				double d2 = p2e1.distanceTo(p2e2);
				
				if (d2 - d1 >= getThreshold(PredicateNames.DISTANCE_INCREASING, pair[0])) {
					frame.add(new Predicate(PredicateNames.DISTANCE_INCREASING, pair[0].getPredArg(), pair[1].getPredArg()));
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	/*
	 * Extract all the DISTANCE_DECREASING_TO_STATIONARY predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistanceDecreasingToStationary() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				Point3D p1e1 = pair[0].getPoint(i);
				Point3D p1e2 = pair[1].getPoint(i);
				Point3D p2e1 = pair[0].getPoint(i+1);
				
				if (p1e1 == null || p1e2 == null || p2e1 == null) {
					continue;
				}
				
				double d1 = p1e1.distanceTo(p1e2);
				double d2 = p2e1.distanceTo(p1e2);
				
				if (d1 - d2 >= getThreshold(PredicateNames.DISTANCE_DECREASING_TO_STATIONARY, pair[0])) {
					frame.add(new Predicate(PredicateNames.DISTANCE_DECREASING_TO_STATIONARY, pair[0].getPredArg(), pair[1].getPredArg()));
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
		
	/*
	 * Extract all the DISTANCE_INCREASING_TO_STATIONARY predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistanceIncreasingToStationary() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				Point3D p1e1 = pair[0].getPoint(i);
				Point3D p1e2 = pair[1].getPoint(i);
				Point3D p2e1 = pair[0].getPoint(i+1);
				
				if (p1e1 == null || p1e2 == null || p2e1 == null) {
					continue;
				}
				
				double d1 = p1e1.distanceTo(p1e2);
				double d2 = p2e1.distanceTo(p1e2);
				
				if (d2 - d1 >= getThreshold(PredicateNames.DISTANCE_INCREASING_TO_STATIONARY, pair[0])) {
					frame.add(new Predicate(PredicateNames.DISTANCE_INCREASING_TO_STATIONARY, pair[0].getPredArg(), pair[1].getPredArg()));
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}	
	
	/*
	 * Extract all the DISTANCE_STATIONARY predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistanceStationary() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				Point3D p1e1 = pair[0].getPoint(i);
				Point3D p1e2 = pair[1].getPoint(i);
				Point3D p2e1 = pair[0].getPoint(i+1);
				Point3D p2e2 = pair[1].getPoint(i+1);
				
				if (p1e1 == null || p1e2 == null || p2e1 == null || p2e2 == null) {
					continue;
				}
				
				double d1 = p1e1.distanceTo(p1e2);
				double d2 = p2e1.distanceTo(p2e2);

				if (Math.abs(d2-d1) < getThreshold(PredicateNames.DISTANCE_STABLE, pair[0])) {
					frame.add(new Predicate(PredicateNames.DISTANCE_STABLE, pair[0].getPredArg(), pair[1].getPredArg()));
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	public Vector<Vector<Predicate>> extractTrajectoryStable() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		retval.add(new Vector<Predicate>());
		for (int i = 1; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e : scene.getEntities()) {
				Point3D p1 = e.getPoint(i-1);
				Point3D p2 = e.getPoint(i);
				Point3D p3 = e.getPoint(i+1);
				
				if (p1 != null && p2 != null && p3 != null) {
					Vec3D v1 = new Vec3D(p1, p2);
					Vec3D v2 = new Vec3D(p2, p3);
					double angle = v1.angle(v2);
					if (Math.abs(angle) < getThreshold(PredicateNames.TRAJECTORY_STABLE, e)) {
						frame.add(new Predicate(PredicateNames.TRAJECTORY_STABLE, e.getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	public Vector<Vector<Predicate>> extractSpeedIncreasing() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		retval.add(new Vector<Predicate>());
		for (int i = 1; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e : scene.getEntities()) {
				Point3D p1 = e.getPoint(i-1);
				Point3D p2 = e.getPoint(i);
				Point3D p3 = e.getPoint(i+1);
				
				if (p1 != null && p2 != null && p3 != null) {
					double accel = p2.distanceTo(p3) - p1.distanceTo(p2);
					if (accel > getThreshold(PredicateNames.SPEED_INCREASING, e)) {
						frame.add(new Predicate(PredicateNames.SPEED_INCREASING, e.getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	public Vector<Vector<Predicate>> extractSpeedDecreasing() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		retval.add(new Vector<Predicate>());
		for (int i = 1; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e : scene.getEntities()) {
				Point3D p1 = e.getPoint(i-1);
				Point3D p2 = e.getPoint(i);
				Point3D p3 = e.getPoint(i+1);
				
				if (p1 != null && p2 != null && p3 != null) {
					double accel = p1.distanceTo(p2) - p2.distanceTo(p3);
					if (accel > getThreshold(PredicateNames.SPEED_DECREASING, e)) {
						frame.add(new Predicate(PredicateNames.SPEED_DECREASING, e.getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}
	
	// returns a Vector of pairs of SceneEntity
	// it can be used for the detection of predicates like distance-increasing
	// in which we have to compare a SceneEntity to every other one
	private Vector<SceneEntity[]> pairPermutations() {
		Vector<SceneEntity[]> pairs = new Vector<SceneEntity[]>();
		Iterator<SceneEntity> i1 = scene.getEntities().iterator();
		while (i1.hasNext()) {
			SceneEntity e1 = i1.next();
			Iterator<SceneEntity> i2 = scene.getEntities().iterator();
			while (i2.hasNext()) {
				SceneEntity e2 = i2.next();
				if (!e1.equals(e2)) {
					SceneEntity[] pair = new SceneEntity[2];
					pair[0] = e1;
					pair[1] = e2;
					pairs.add(pair);
				}
			}
		}
		return pairs;
	}
}
